import os, sys
import numpy as np
import h5py
import sigpy as sp
import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"] = (10,10)
import sigpy.plot as pl
import pathlib
path = os.environ["TOOLBOX_PATH"] + "/python/"
sys.path.append(path)
import bart
import tqdm
from tqdm import trange
import shutil
import copy





mask = np.load('../klf_project/dataset/mask.npy')

for fid_filename in os.listdir('../klf_project/dataset/fid'):
    fid_path = '../klf_project/dataset/fid/'+fid_filename
    if not os.path.exists('dataset/'+fid_filename.split('_')[0]):
        os.mkdir('dataset/'+fid_filename.split('_')[0])
    if not os.path.exists('dataset/'+fid_filename.split('_')[0]+'/Train'):
        os.mkdir('dataset/'+fid_filename.split('_')[0]+'/Train')
    if not os.path.exists('dataset/'+fid_filename.split('_')[0]+'/Val'):
        os.mkdir('dataset/'+fid_filename.split('_')[0]+'/Val')
    fid_data = np.load(fid_path)
    
    

    # every patient 
    for index in trange (fid_data.shape[0]) :
        if index < fid_data.shape[0]*0.8:
            h5file = 'dataset/'+fid_filename.split('_')[0]+'/Train/'+str(index)+'.h5'
        else:
            h5file = 'dataset/'+fid_filename.split('_')[0]+'/Val/'+str(index)+'.h5'

        hf = h5py.File(h5file,'a')
        raw_data = fid_data[index]
        
        kspaces = []
        esp_maps = []
        # every slice 
        for slice in range(raw_data.shape[0]):
            

            cropped_size = (raw_data.shape[2] - 256 )//2
            

            kspace_slice = raw_data[slice,: ,cropped_size:raw_data.shape[2]-cropped_size, :]

            # kspace_slice = raw_data[slice,: :, :]

            ksp_slice = np.moveaxis(kspace_slice,0,-1)
            maps = bart.bart(1,'ecalib -m 1 -a -d 0 -g  -r 24',ksp_slice[None])
            esp_maps.append(np.moveaxis(maps[0],-1,0))
            kspaces.append(kspace_slice)

        kspaces = np.stack(kspaces,0)
        esp_maps = np.stack(esp_maps,0)
        print('Saving the kspaces of size:', kspaces.shape)
        print('Saving the espirit_maps of size:', esp_maps.shape)
        hf.create_dataset('kspace', data=kspaces)
        hf.create_dataset('esp_maps', data=esp_maps)
        hf.close()

    